import numpy as np
from pulp import pulp, LpMaximize, LpInteger, PULP_CBC_CMD


def max_function(Y, X, n, m):
    s = 0
    for i in range(n):
        for j in range(m):
            s += Y[i][j] + X[j][i]
    return s


def generate_hours(coefficients, user_constraints, prediction_constraints):
    Y = np.array(coefficients)
    n, m = Y.shape

    prob = pulp.LpProblem("Maximize_Trace", LpMaximize)

    # Определяем матрицу X с часами
    # строки - модальности
    # столбцы - врачи
    X = pulp.LpVariable.matrix("X", (range(m), range(n)), lowBound=0, cat=LpInteger)

    # Максимизируем след матрицы Y * X можно уменьшить количество вычислений, так как для вычисления следа не требуется полное перемножение

    prob += max_function(Y, X, n, m)

    # задаем сколько часов было заложено на предсказанную модальность
    for i in range(m):
        prob += pulp.lpSum([X[i][j] for j in range(n)]) <= prediction_constraints[i]

    # задаем сколько часов у врачей предполагается по графику
    for i in range(n):
        prob += pulp.lpSum([X[j][i] for j in range(m)]) <= user_constraints[i]

    prob.solve(PULP_CBC_CMD(msg=False))

    X_optimized = np.array([[X[i][j].varValue for j in range(n)] for i in range(m)])

    return X_optimized
